package com.aizuda.easyManagerTool.service.gpt;

import cn.hutool.json.JSONUtil;
import com.aizuda.easyManagerTool.domain.vo.socket.SocketMessageVO;
import com.aizuda.easyManagerTool.service.gpt.impl.XFAiGPTServiceImpl;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.plexpt.chatgpt.entity.chat.Message;
import lombok.AllArgsConstructor;
import lombok.Data;
import lombok.NoArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import okhttp3.Response;
import okhttp3.WebSocket;
import okhttp3.WebSocketListener;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;

import java.io.IOException;
import java.util.List;
import java.util.function.Consumer;
import java.util.stream.Collectors;

@Slf4j
public class WebSocketStreamListener  extends WebSocketListener {


    final WebSocketSession webSocketSession;
    SocketMessageVO socketMessageVO = new SocketMessageVO();
    final static String TYPE = "chartGPT";
    final XFAiGPTServiceImpl.XFObject xfObject;
    final MessageModel model;

    StringBuffer sb = new StringBuffer();
    ObjectMapper objectMapper = new ObjectMapper();

    protected Consumer<String> onComplate = (s) -> {
    };

    public WebSocketStreamListener(WebSocketSession webSocketSession, XFAiGPTServiceImpl.XFObject xfObject, MessageModel model) {
        this.webSocketSession = webSocketSession;
        this.xfObject = xfObject;
        this.model = model;
    }

    @Override
    public void onMessage(@NotNull WebSocket webSocket, @NotNull String text) {
        // 解析text
        try {
            JsonParse jsonParse = JSONUtil.toBean(text, JsonParse.class);
            Header header = jsonParse.header;
            // 消息发生错误
            if (header.code != 0) {
                send(jsonParse.header.message, "assistant");
                sb = null;
                return;
            }
            List<Text> textList = jsonParse.payload.choices.text;
            String collect = textList.stream().map(i -> i.content).collect(Collectors.joining());
            collect = collect.replaceAll("\\n","\n");
            sb.append(collect);
            send(collect, textList.get(0).getRole());
            // 结束
            if (jsonParse.header.status == 2) {
                // 可以关闭连接，释放资源
                this.onComplate.accept(sb.toString());
            }
        }catch (Exception e){
            e.printStackTrace();
            send(e.getMessage(),"assistant");
        }finally {
            webSocket.close(1000, "Closing WebSocket connection");
        }
    }

    @Override
    public void onOpen(@NotNull WebSocket webSocket, @NotNull Response response) {
        super.onOpen(webSocket, response);
        // 把历史数据发给服务器
        try {
            webSocket.send(objectMapper.writeValueAsString(xfObject));
        }catch (Exception e){
            log.error("发送给讯飞历史消息出错 {}",e.getMessage());
        }
    }

    @Override
    public void onClosed(@NotNull WebSocket webSocket, int code, @NotNull String reason) {
        super.onClosed(webSocket, code, reason);
    }

    @Override
    public void onFailure(@NotNull WebSocket webSocket, @NotNull Throwable t, @Nullable Response response) {
        super.onFailure(webSocket, t, response);
        try {
            if (null != response) {
                int code = response.code();
                log.error("讯飞AI错误 ,code = {} result={}",code,response.body().string());
                if (101 != code) {
                    send("连接失败","assistant");
                }else{
                    send("讯飞AI错误 错误码 = "+code+" 原因 = "+response.body().string(),"assistant");
                }
            }
        } catch (IOException e) {
            log.error("讯飞接收数据报错 {}",e.getMessage());
        }finally {
            webSocket.close(1000, "Closing WebSocket connection");
        }
    }

    public void setOnComplate(Consumer<String> onComplate) {
        this.onComplate = onComplate;
    }

    private void send(String s,String role){
        try {
            Message build = Message.builder().name("ai").role(role).content(s).build();
            // 具体哪个功能
            socketMessageVO.setType(TYPE);
            model.setMsg(build);
            socketMessageVO.setObj(model);
            webSocketSession.sendMessage(new TextMessage(JSONUtil.toJsonStr(socketMessageVO)));
        }catch (Exception e){
            log.error(e.getMessage());
        }
    }


    @AllArgsConstructor
    @NoArgsConstructor
    @Data
    class JsonParse {
        Header header;
        Payload payload;
    }

    @AllArgsConstructor
    @NoArgsConstructor
    @Data
    class Header {

        int code;
        int status;
        String message;
        String sid;
    }

    @AllArgsConstructor
    @NoArgsConstructor
    @Data
    class Payload {
        Choices choices;
    }

    @AllArgsConstructor
    @NoArgsConstructor
    @Data
    class Choices {
        int status;
        int seq;
        List<Text> text;
    }

    @AllArgsConstructor
    @NoArgsConstructor
    @Data
    class Text {
        String role;
        String content;
        int index;
    }

}
